import random
import time
import matplotlib.pyplot as plt


# 1. 排序法：升序排序后取倒数第k个元素
def find_kth_largest_sort(arr, k):
    arr_copy = arr.copy()  # 避免修改原数组
    arr_copy.sort()  # 升序排序
    return arr_copy[len(arr_copy) - k]

# 2. 最小栈法：维护大小为k的升序栈（栈顶最小）
def find_kth_largest_min_stack(arr, k):
    min_stack = []
    for num in arr:
        if len(min_stack) < k:
            min_stack.append(num)
            min_stack.sort()  # 压栈后排序，维持升序（栈顶最小）
        else:
            if num > min_stack[0]:  # 当前元素大于栈顶（最小候选），则替换
                min_stack.pop(0)
                min_stack.append(num)
                min_stack.sort()  # 重新排序维持结构
    return min_stack[0]  # 栈顶为第k大元素

# ----------------------耗时统计（生成实验数据）----------------------
def get_performance_data():
    # 实验报告中的数组规模梯度
    n_list = [1000, 5000, 10000, 20000, 50000, 100000, 150000]
    sort_avg_times = []  # 排序法平均耗时（ms）
    stack_avg_times = []  # 最小栈法平均耗时（ms）
    repeat = 3  # 每组规模重复3次，取平均

    for n in n_list:
        k = max(1, int(n * 0.1))  # k=数组规模的10%（向上取整）
        sort_total = 0.0
        stack_total = 0.0

        # 重复3次测试，累计耗时
        for _ in range(repeat):
            # 生成随机数组（元素范围1~100000）
            arr = [random.randint(1, 100000) for _ in range(n)]

            # 统计排序法耗时
            start = time.time()
            find_kth_largest_sort(arr, k)
            end = time.time()
            sort_total += (end - start) * 1000  # 转换为毫秒

            # 统计最小栈法耗时
            start = time.time()
            find_kth_largest_min_stack(arr, k)
            end = time.time()
            stack_total += (end - start) * 1000

        # 计算平均耗时
        sort_avg_times.append(round(sort_total / repeat, 2))
        stack_avg_times.append(round(stack_total / repeat, 2))

    return n_list, sort_avg_times, stack_avg_times

# ----------------------绘制趋势曲线----------------------
def plot_trend_curve():
    n_list, sort_times, stack_times = get_performance_data()

    # 设置画布样式
    plt.figure(figsize=(12, 6))  # 画布大小（宽12，高6）
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 解决中文显示问题
    plt.rcParams['axes.unicode_minus'] = False

    # 绘制两条趋势线
    plt.plot(n_list, sort_times, 
             label='排序法（O(n log n)）', 
             marker='o',  # 数据点用圆形标记
             color='#1f77b4',  # 蓝色
             linewidth=2,  # 线宽
             markersize=6)  # 标记大小

    plt.plot(n_list, stack_times, 
             label='最小栈法（实际近O(n²)）', 
             marker='s',  # 数据点用方形标记
             color='#ff7f0e',  # 橙色
             linewidth=2,
             markersize=6)

    # 设置坐标轴与标题
    plt.xlabel('数组规模（n）', fontsize=12)
    plt.ylabel('平均运行时间（ms）', fontsize=12)
    plt.title('排序法vs最小栈法 查找第K大元素的时间复杂度趋势', fontsize=14, pad=20)
    plt.legend(fontsize=11)  # 图例
    plt.grid(True, linestyle='--', alpha=0.5)  # 网格线（虚线，半透明）

    # 保存图片（可直接在当前文件夹查看）
    plt.savefig('kth_largest_trend_curve.png', dpi=300, bbox_inches='tight')
    plt.show()  # 显示图片

# ----------------------执行代码----------------------
if __name__ == "__main__":
    plot_trend_curve()